import json
import numpy as np
import faiss
from text2vec import SentenceModel

# 1. 读取题库 JSON 文件
with open("back.json", 'r', encoding='utf-8') as f:
    data = json.load(f)

# 2. 加载中文文本向量模型
model = SentenceModel("shibing624/text2vec-base-chinese")

# 3. 为每个问题生成向量，并存入数据中
for q in data:
    q["vector"] = model.encode(q["question"])

# 4. 构建向量矩阵（用于 FAISS）
vectors = np.array([q["vector"] for q in data], dtype='float32')

# 5. 构建 FAISS 索引
index = faiss.IndexFlatL2(vectors.shape[1])  # 使用 L2 距离, IndexFlatL2平面索引
index.add(vectors)

# 6. 保存向量矩阵和索引到磁盘
np.save("back_vectors.npy", vectors)
faiss.write_index(index, "back_questions.index")

# 7. 保存带向量的题库为 JSON（注意需将 ndarray 转换为 list）
for q in data:
    if isinstance(q.get("vector"), np.ndarray):
        q["vector"] = q["vector"].tolist()

with open("back_questions_with_vectors.json", "w", encoding="utf-8") as f:
    json.dump(data, f, ensure_ascii=False, indent=2)

print("处理完成：向量和索引已保存。")
